import os
from pathlib import Path
import re

def extract_image_info_and_save(root_dir: str, output_file: str = "extracted_images.txt",output_file_dict: str = "extracted_images_dict.txt"):
    """
    遍历指定目录及其子目录，提取符合特定模式的图片文件名信息，并保存到文本文件中。
    
    假设图片文件名格式为: <label> <description>_<random1>_<random2>.<ext>
    例如: 56 awarded_7076_2078379.png
    提取后格式: relative_path_to_image label description
    例如: images/train_word_1.png 56 awarded

    Args:
        root_dir (str): 要遍历的根目录路径
        output_file (str): 输出的文本文件名，默认为 'extracted_images.txt'
    """
    
    # 定义支持的图片扩展名
    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'}

    extracted_data = []
    extracted_data_dict = {}
    root_path = Path(root_dir)
    
    if not root_path.exists():
        raise FileNotFoundError(f"指定的目录不存在: {root_dir}")
    
    if not root_path.is_dir():
        raise NotADirectoryError(f"指定的路径不是一个目录: {root_dir}")
    
    # 遍历目录及其子目录
    for file_path in root_path.rglob('*'):
        if file_path.is_file() and file_path.suffix.lower() in image_extensions:

            relative_path = file_path.relative_to(root_path)
            # 使用正斜杠作为路径分隔符以确保跨平台兼容性
            formatted_path = str(relative_path).replace('\\', '/')

            # 提取标签和描述部分（第一组）
            label_description = formatted_path.split("/")[-1].split("_")[0].strip()
            for i_ in label_description.replace(' ', ''):
                i=str(i_)
                if i not in extracted_data_dict:
                    extracted_data_dict[i]=0
                extracted_data_dict[i]+=1
                print(i)

            # 获取相对于根目录的路径
            try:
                # 构建输出行
                output_line = f"{formatted_path} {label_description}"
                extracted_data.append(output_line)
                
                # print(f"已处理: {output_line}")
                
            except ValueError:
                # 如果文件不在根目录下（理论上不会发生，因为使用了rglob）
                print(f"警告: 文件不在根目录内，跳过: {file_path}")
    
    # 将结果写入文件
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            for line in extracted_data:
                f.write(line + '\n')
        print(f"\n成功提取 {len(extracted_data)} 个图片信息。")
        print(f"结果已保存到: {os.path.abspath(output_file)}")
        
        with open(output_file_dict, 'w', encoding='utf-8') as f:
            for line in extracted_data_dict:
                f.write(line + '\n')
        print(f"\n成功提取 {extracted_data_dict}")
        
    except IOError as e:
        print(f"写入文件时发生错误: {e}")
        raise

# -------------------------- 使用示例 --------------------------
    
if __name__ == "__main__":
    # 请修改这里的路径为您的实际目录
    DIRECTORY_TO_SCAN = "/root/PaddleOCR/dataset/en_rec_3W_res"  # <-- 修改这里
    OUTPUT_FILENAME = "/root/PaddleOCR/dataset/en_rec_3W_res/en_word.txt"      # <-- 可选：修改输出文件名
    OUTPUT_FILENAME_dict = "/root/PaddleOCR/dataset/en_rec_3W_res/en_word_dict.txt"      # <-- 可选：修改输出文件名
    
    try:
        extract_image_info_and_save(DIRECTORY_TO_SCAN, OUTPUT_FILENAME,OUTPUT_FILENAME_dict)
    except (FileNotFoundError, NotADirectoryError) as e:
        print(f"错误: {e}")
        print("请检查您设置的目录路径是否正确。")
    except Exception as e:
        print(f"发生未预期的错误: {e}")